# -*- coding: utf-8 -*-
# @Time    : 2021/4/19
# @File    : stylizer.py
import os

from PIL import Image
from torch.autograd import Variable
from torchvision.utils import save_image

from models import TransformerNet
from tool.utils import *

# style_name = 'moisac/moisac12000.pth'
# style_name = 'udnie/udnie33000.pth'
style_name = 'starry/starry50000.pth'

args = {
	'image_path': 'images/content/1.jpg',
	'checkpoint_model': f'./checkpoints/{style_name}',
}

if __name__ == "__main__":
	os.makedirs("images/outputs", exist_ok=True)
	# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
	device = torch.device("cpu")

	transform = style_transform(None)

	# Define model and load model checkpoint
	transformer = TransformerNet().to(device)
	transformer.load_state_dict(torch.load(args['checkpoint_model'])['model'])
	transformer.eval()

	# Prepare input
	image_tensor = Variable(transform(Image.open(args['image_path']).convert('RGB'))).to(device)
	image_tensor = image_tensor.unsqueeze(0)

	# Stylize image
	with torch.no_grad():
		stylized_image = denormalize(transformer(image_tensor)).cpu()

	# Save image
	fn = args['image_path'].split("/")[-1]
	# style_name =
	save_image(stylized_image, f"images/outputs/{style_name.split('/')[-1][:-4]}-{fn}")
